from preprocess.data_utils import file_conversion, parse_files, prepare_vectors, split_to_chunk
from preprocess.file_paser.img_parser.img_ocr import image_ocr

dataset_dir = r'D:\Users\JHC258\projects\文本检索\preprocess\dataset'
parse_out_dir = r'D:\Users\JHC258\projects\文本检索\preprocess\parse_out'
ofd_unzip_dir = r'D:\Users\JHC258\projects\文本检索\preprocess\ofd_unzipped'
vec_data_dir = r'D:\Users\JHC258\projects\文本检索\vec_db_data'
image_saved_dir = r'D:\Users\JHC258\projects\文本检索\preprocess\images'
block_max_len = 400  # 块的最大长度，取值100~500
store_batch_size = 3000  # 向量化的数据按此大小划分存储以减少内存压力，建议取3000


def insert_data():
    from milvus.milvus_op import MilvusOP
    from milvus.utils import traverse_dir
    from tqdm import tqdm

    milvus_op = MilvusOP(db_name='state_vector_db', collection_name="hybrid3")
    print('Inserting data ...')
    path_iter = traverse_dir(path=r'../vec_db_data', tails=['pkl'])
    for f_pth, f_name in tqdm(path_iter, desc='正在插入Milvus'):
        milvus_op.insert_data(f_pth)
    print('Completed!.')


def preprocess_files():
    # 1. 使用dataset_path下文件，其中，doc、pdf转docx文件，xls转xlsx文件
    file_conversion(dataset_dir)

    # 2. 文件转换后，将文件内容解析下来。结果以json保存到parse_out
    parse_files(dataset_dir, parse_out_dir, ofd_unzip_dir, image_saved_dir, block_max_len)

    # 3.提取图片中的文字
    image_ocr(parse_out_dir, image_saved_dir)

    # 4. 按照split_size将文件内容划为不同块。结果以json保存到prepared_data_dir
    split_to_chunk(parse_out_dir, vec_data_dir, block_max_len)

    # 5. 向量化不同的块。结果以pkl保存到prepared_data_dir
    prepare_vectors(vec_data_dir, block_max_len, store_batch_size)


if __name__ == "__main__":
    # 2. 文件转换后，将文件内容解析下来。结果以json保存到parse_out
    parse_files(dataset_dir, parse_out_dir, ofd_unzip_dir, image_saved_dir, block_max_len)

    # 3.提取图片中的文字
    image_ocr(parse_out_dir, image_saved_dir)

    # 4. 按照split_size将文件内容划为不同块。结果以json保存到prepared_data_dir
    split_to_chunk(parse_out_dir, vec_data_dir, block_max_len)

    # 5. 向量化不同的块。结果以pkl保存到prepared_data_dir
    prepare_vectors(vec_data_dir, block_max_len, store_batch_size)

    insert_data()
